/*
 * Copyright 2010 Bo Zhou<Bo.Schwarzstein@gmail.com>
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 
		http://www.apache.org/licenses/LICENSE-2.0 
 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "MeshSmoothing.h"

#include <cassert>
#include <iostream>
#include <set>

#include <maya/MPointArray.h>
#include <maya/MIntArray.h>

MeshSmoothing::MeshSmoothing(const MObject& MeshObject) : mMesh(MeshObject), mUsePaintedVertices(false)
{
}

MeshSmoothing::~MeshSmoothing()
{
}

void MeshSmoothing::UsePaintedVertices(const bool Use)
{
	mUsePaintedVertices = Use;
}

bool MeshSmoothing::IsUsingPaintedVertices() const
{
	return mUsePaintedVertices;
}

void MeshSmoothing::SetPaintedColor(const MColor& PaintedColor)
{
	mPaintedColor = PaintedColor;
}

const MColor& MeshSmoothing::GetPaintedColor() const
{
	return mPaintedColor;
}

bool MeshSmoothing::Prepare()
{
	try
	{
		const MFnMesh& Mesh = mMesh;

		const int PolygonCount = Mesh.numPolygons();
		const int VertexCount =  Mesh.numVertices();

		// Calculate a table which from each vertex to the polygon id from at which refered the vertex.
		//
		MIntArray Counts, Indices;
		Mesh.getVertices(Counts, Indices);

		std::valarray<int> VertexReferedTimes(0,VertexCount);
		int Offset = 0;
		for(unsigned int p=0; p<Counts.length(); ++p)
		{
			const int Count = Counts[p];
			for(int i=0; i<Count; ++i)
			{
				int j = Offset + i;

				const int Index = Indices[j];
				VertexReferedTimes[Index] += 1;
			}
			Offset += Count;
		}

		const int Arity = VertexReferedTimes.max() + 1;
		mVertexSharedPolygons.resize(Arity*VertexCount);
		std::fill_n(&mVertexSharedPolygons[0],mVertexSharedPolygons.size(),-1);
		std::valarray<int> Offsets(0,VertexCount);

		Offset = 0;
		for(unsigned int p=0; p<Counts.length(); ++p)
		{
			const int Count = Counts[p];
			for(int i=0; i<Count; ++i)
			{
				int j = Offset + i;

				const int Index = Indices[j];

				const int InsertPos = Index*Arity + Offsets[Index];
				mVertexSharedPolygons[InsertPos] = p;

				Offsets[Index] += 1;
			}
			Offset += Count;
		}
		
		mMaxVertexSharedTime = Arity - 1;

		// Get the all painted vertex ids.
		//
		if( IsUsingPaintedVertices() )
		{
			MColor PaintColor( GetPaintedColor() );

			MColorArray VertexColors;
			mMesh.getVertexColors(VertexColors);

			size_t PaintedVerticesSize = 0;
			for(unsigned int i=0; i<VertexColors.length(); ++i)
			{
				const MColor VC(VertexColors[i]);
				if( VC == PaintColor )
					++PaintedVerticesSize;
			}
			mPaintedVertices.resize( PaintedVerticesSize, -1 );

			for(unsigned int i=0, j=0; i<VertexColors.length(); ++i)
			{
				const MColor VC(VertexColors[i]);
				if( VC == PaintColor )
				{
					mPaintedVertices[j++] = (int)i;
				}
			}
		}

		return true;
	}catch(const std::exception& e)
	{
		std::cerr<<e.what()<<std::endl;
		return false;
	}catch(...)
	{
		return false;
	}
}

bool MeshSmoothing::Smooth()
{
	// Do smoothing.
	// 
	const int VertexCount = mMesh.numVertices();
	MPointArray NewPoints(VertexCount);

	std::valarray<int> VertexIds( VertexCount );
	for(size_t i=0; i<VertexIds.size(); ++i)
	{
		VertexIds[i] = (int)i;
	}

	if( IsUsingPaintedVertices() )
	{
		VertexIds = mPaintedVertices;
	}

	for(size_t i=0; i<VertexIds.size(); ++i)
	{
		const int vid = VertexIds[i];
		assert( vid >= 0 );

		std::set<int> Neighbors;

		// Get neighborhood of indices
		// 
		int* p = &mVertexSharedPolygons[vid*(mMaxVertexSharedTime+1)];
		while( *p > -1 )
		{
			const int Polygon = *p;
			MIntArray IndicesOnPolygon;
			mMesh.getPolygonVertices(Polygon, IndicesOnPolygon);
			for(unsigned int j=0; j<IndicesOnPolygon.length(); ++j)
			{
				Neighbors.insert( IndicesOnPolygon[j] );
			}
			++p;
		}

		std::set<int>::iterator itr = Neighbors.find(vid);
		assert( itr != Neighbors.end() );
		Neighbors.erase(itr);

		MPoint AverageP;
		for(itr=Neighbors.begin(); itr!=Neighbors.end(); ++itr)
		{
			MPoint P;
			mMesh.getPoint(*itr, P);
			AverageP += P;
		}
		double Weight = 1.0 / (double)Neighbors.size();
		AverageP.x *= Weight;
		AverageP.y *= Weight;
		AverageP.z *= Weight;

		NewPoints[vid] = AverageP;
	}

	mMesh.setPoints(NewPoints);
	
	NewPoints.clear();
	NewPoints.setLength(0);
	return true;
}
